from backup.rsa_core import *


# # --- 主程序示例 ---
# if __name__ == '__main__':
#     # 设置密钥位数 (例如 2024 位，p 和 q 大约是 1012 位)
#     # 为了演示快速，可以用小一点的位数，比如 64 或 128
#     # 注意：低于 1024 位的 RSA 在实践中是不安全的
#     # key_bits = 64 # 用于快速测试
#     key_bits = 2048 # 按要求
#     # key_bits = 2048 # 更标准的位数

#     print(f"开始生成 {key_bits} 位 RSA 密钥对...\n")
#     public_key, private_key = generate_keypair(bits=key_bits)

#     print("\n--- 密钥生成完毕 ---")
#     print(f"公钥 (n, e):")
#     print(f"n = {public_key[0]}")
#     print(f"e = {public_key[1]}\n")

#     print(f"私钥 (n, d):")
#     print(f"n = {private_key[0]} (与公钥中的n相同)")
#     print(f"d = {private_key[1]}\n")

#     # 要加密的消息
#     message_str = "你好，RSA！这是一个测试消息。"
#     message_bytes = message_str.encode('utf-8')
#     print(f"原始消息 (字符串): {message_str}")
#     print(f"原始消息 (bytes): {message_bytes}\n")

#     # 加密
#     print("--- 开始加密 ---")
#     try:
#         ciphertext = encrypt(public_key, message_bytes)
#         print(f"密文 (整数): {ciphertext}\n")

#         # 解密
#         print("--- 开始解密 ---")
#         decrypted_bytes = decrypt(private_key, ciphertext)
#         decrypted_message_str = decrypted_bytes.decode('utf-8', errors='replace') # 使用replace处理可能存在的解码错误

#         print(f"解密后的消息 (bytes): {decrypted_bytes}")
#         print(f"解密后的消息 (字符串): {decrypted_message_str}\n")

#         # 验证
#         if message_bytes == decrypted_bytes:
#             print("成功！加密和解密结果一致。")
#         else:
#             print("失败！加密和解密结果不一致。")
#             print(f"原始: {message_bytes}")
#             print(f"解密后: {decrypted_bytes}")

#     except ValueError as e:
#         print(f"发生错误: {e}")
#         print("这通常意味着消息对于当前密钥大小来说太长了。")
#         print("对于教科书RSA，消息转换的整数 M 必须小于模数 N。")
#         print(f"消息转换后的整数大小: {math_utils.bytes_to_int(message_bytes).bit_length()} bits")
#         print(f"模数 N 的大小: {public_key[0].bit_length()} bits")
#         print("在实际应用中，会使用填充方案（如OAEP）来处理任意长度的消息，并增强安全性。")
        
        
        
        
        
        
        
        

# --- 主程序示例 ---
if __name__ == '__main__':
    # key_bits = 256 # 快速测试，注意消息长度限制会更严格
    key_bits = 4096 # 标准位数
    # key_bits = 2024 # 按之前要求

    print(f"开始生成 {key_bits} 位 RSA 密钥对...\n")
    public_key, private_key = generate_keypair(bits=key_bits)

    print("\n--- 密钥生成完毕 ---")
    print(f"公钥 (n, e): n={public_key[0]}, e={public_key[1]}")
    # print(f"私钥 (n, d): n={private_key[0]}, d={private_key[1]}") # 私钥d通常不显示

    # 要加密的消息
    message_str_short = "Hello OAEP!"
    message_bytes_short = message_str_short.encode('utf-8')

    # 对于2048位密钥 (k_bytes=256), SHA-256 (h_len=32)
    # max_msg_len = 256 - 2*32 - 2 = 256 - 64 - 2 = 190 bytes
    # 一个较长的消息，但仍在限制内
    message_str_long = "这是一个用于测试RSA-OAEP的相对较长的消息。希望它能正常工作！" * 2
    message_bytes_long = message_str_long.encode('utf-8')
    print(f"长消息字节长度: {len(message_bytes_long)}")

    # 标签 (可选)
    label = b"MyApplication-Context-Specific-Label"

    messages_to_test = {
        "短消息": message_bytes_short,
        "长消息": message_bytes_long,
    }

    for msg_name, message_bytes in messages_to_test.items():
        print(f"\n--- 测试: {msg_name} ---")
        print(f"原始消息 (bytes): {message_bytes[:60]}..." if len(message_bytes) > 60 else message_bytes)

        # 加密
        print("--- 开始加密 (OAEP) ---")
        try:
            ciphertext = encrypt_oaep(public_key, message_bytes, label_bytes=label)
            print(f"密文 (整数): {ciphertext}\n")

            # 解密
            print("--- 开始解密 (OAEP) ---")
            decrypted_bytes = decrypt_oaep(private_key, ciphertext, label_bytes=label)
            # decrypted_message_str = decrypted_bytes.decode('utf-8', errors='replace')

            # print(f"解密后的消息 (bytes): {decrypted_bytes[:60]}..." if len(decrypted_bytes) > 60 else decrypted_bytes)
            # print(f"解密后的消息 (字符串): {decrypted_message_str[:60]}..." if len(decrypted_message_str) > 60 else decrypted_message_str)

            # 验证
            if message_bytes == decrypted_bytes:
                print(f"成功！{msg_name} 加密和解密结果一致。")
            else:
                print(f"失败！{msg_name} 加密和解密结果不一致。")
                print(f"原始: {message_bytes}")
                print(f"解密后: {decrypted_bytes}")

        except ValueError as e:
            print(f"发生错误: {e}")
            n_val = public_key[0]
            k_bytes_val = (n_val.bit_length() + 7) // 8
            h_len_val = DEFAULT_HASH_FUNC().digest_size
            max_len = k_bytes_val - 2 * h_len_val - 2
            print(f"  (当前密钥长度 {key_bits} bits -> k_bytes={k_bytes_val}, hash_len={h_len_val} bytes)")
            print(f"  (OAEP最大消息长度: {max_len} bytes, 当前消息长度: {len(message_bytes)} bytes)")


    # 测试消息过长的情况
    print("\n--- 测试消息过长 ---")
    n_val = public_key[0]
    k_bytes_val = (n_val.bit_length() + 7) // 8
    h_len_val = DEFAULT_HASH_FUNC().digest_size
    max_len = k_bytes_val - 2 * h_len_val - 2
    too_long_message = b'A' * (max_len + 1)
    try:
        encrypt_oaep(public_key, too_long_message, label_bytes=label)
    except ValueError as e:
        print(f"预期错误 (消息过长): {e}")
        assert "Message too long" in str(e)

    # 测试解密错误 (例如，使用错误的标签解密)
    print("\n--- 测试解密错误 (错误标签) ---")
    ciphertext_for_wrong_label = encrypt_oaep(public_key, message_bytes_short, label_bytes=label)
    wrong_label = b"ThisIsTheWrongLabel"
    try:
        decrypt_oaep(private_key, ciphertext_for_wrong_label, label_bytes=wrong_label)
    except ValueError as e:
        print(f"预期错误 (错误标签导致解密失败): {e}")
        assert "Decryption error" in str(e) # 应该是通用错误

    print("\n--- 测试解密错误 (损坏的密文 - 修改一位) ---")
    original_ciphertext = encrypt_oaep(public_key, message_bytes_short, label_bytes=label)
    # 将密文整数稍微改动一点，很可能会导致OAEP解码失败
    tampered_ciphertext = original_ciphertext + 1 
    # 或者如果密文很小，可以尝试直接修改转换后的字节串 EM
    # m_enc_int = math_utils.power(original_ciphertext, private_key[1], public_key[0])
    # em_bytes = math_utils.int_to_bytes(m_enc_int, min_length=k_bytes_val)
    # if len(em_bytes) > 5:
    #    tampered_em_list = list(em_bytes)
    #    tampered_em_list[5] = (tampered_em_list[5] + 1) % 256 # 修改EM的一个字节
    #    tampered_em_bytes = bytes(tampered_em_list)
    #    tampered_ciphertext_from_em = math_utils.power(math_utils.bytes_to_int(tampered_em_bytes), public_key[1], public_key[0])
    # else:
    #    tampered_ciphertext_from_em = original_ciphertext + 1 # 简单篡改

    try:
        decrypt_oaep(private_key, tampered_ciphertext, label_bytes=label)
        print("意外：篡改的密文解密成功了（极小概率，或者篡改方式无效）")
    except ValueError as e:
        print(f"预期错误 (损坏的密文导致解密失败): {e}")
        assert "Decryption error" in str(e)